package io.vertx.example.verticle;

import io.vertx.core.AbstractVerticle;
import io.vertx.core.Handler;
import io.vertx.core.buffer.Buffer;
import io.vertx.core.net.NetServer;
import io.vertx.core.net.NetSocket;
import io.vertx.core.streams.Pump;
import io.vertx.example.bean.DirectConnectBean;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong;

import io.vertx.example.util.MultiPumpImpl;
import lombok.extern.slf4j.Slf4j;

/**
 * 转发请求工具.
 * @author zhouwenqing
 * @date 2018/12/29
 */
@Slf4j
public class ForwardConnectVerticle extends AbstractVerticle {

    private static final AtomicLong conId = new AtomicLong(0L);

    private DirectConnectBean connectBean;
    private NetServer netServer;


    public ForwardConnectVerticle(DirectConnectBean connectBean) {
        this.connectBean = connectBean;
    }


    private void connectToCopyAddr(Long myId, NetSocket netSocket, NetSocket mainTargetNetSocket,
        List<NetSocket> copyConnectedNetSockets,
        AtomicInteger connectingNum, DirectConnectBean.TargetAddr copyAddr) {

        vertx.createNetClient().connect(copyAddr.getPort(), copyAddr.getHost(), conRs -> {
            connectingNum.addAndGet(-1);

            if (conRs.succeeded()) {
                copyConnectedNetSockets.add(conRs.result());
            } else {
                log.info(
                    "接收到本地连接关闭事件,conId:{}, clientPort:{},fromHost:{},fromPost:{}, remoteHost:{}, remotePort:{}",
                    myId,
                    connectBean.getClientPort(),
                    netSocket.remoteAddress().host(), netSocket.remoteAddress().port(),
                    connectBean.getHost(), connectBean.getPort());
            }

            if (connectingNum.get() <= 0) {
                startPump(myId, netSocket, mainTargetNetSocket, copyConnectedNetSockets);
            }
        });
    }

    private void startPump(Long myId, NetSocket netSocket, NetSocket mainTargetNetSocket,
        List<NetSocket> copyConnectedNetSockets) {
        String fromHost = netSocket.remoteAddress().host();
        int fromPort = netSocket.remoteAddress().port();
        List<NetSocket> targetWriteSocket = new ArrayList<>();
        targetWriteSocket.add(mainTargetNetSocket);
        targetWriteSocket.addAll(copyConnectedNetSockets);
        Pump writePump = Pump.pump(mainTargetNetSocket, netSocket).start();
        Pump readPump = new MultiPumpImpl(netSocket, targetWriteSocket).start();

        netSocket.resume();


        // 丢弃拷贝目标返回的数据
        Handler<Buffer> emptyHandler = (buf)->{};
        copyConnectedNetSockets.forEach(ci->ci.handler(emptyHandler));

        netSocket.closeHandler(v -> {
            log.info(
                "接收到本地连接关闭事件,conId:{}, clientPort:{},fromHost:{},fromPost:{}, remoteHost:{}, remotePort:{}",
                myId,
                connectBean.getClientPort(),
                fromHost, fromPort,
                connectBean.getHost(), connectBean.getPort());
            readPump.stop();
            writePump.stop();
            mainTargetNetSocket.close();
            copyConnectedNetSockets.forEach(ns -> ns.close());
        });
        mainTargetNetSocket.closeHandler(v -> {
            log.info(
                "接收到远程连接关闭事件,conId:{}, clientPort:{}, fromHost:{},fromPost:{},remoteHost:{}, remotePort:{}",
                myId,
                connectBean.getClientPort(),
                fromHost, fromPort,
                connectBean.getHost(), connectBean.getPort());
            readPump.stop();
            writePump.stop();
            netSocket.close();
            copyConnectedNetSockets.forEach(ns -> ns.close());
        });
    }

    @Override
    public void start() throws Exception {
        log.info("启动监听, listenHost:{}, listenPort:{}, toHost:{}, toPort:{}",
            connectBean.getListenHost(), connectBean.getClientPort(),
            connectBean.getHost(), connectBean.getPort());
        netServer = vertx.createNetServer().connectHandler(netSocket -> {
            netSocket.pause();
            String fromHost = netSocket.remoteAddress().host();
            int fromPort = netSocket.remoteAddress().port();
            Long myId = conId.incrementAndGet();
            vertx.createNetClient().connect(connectBean.getPort(), connectBean.getHost(), conRs -> {
                if (!conRs.succeeded()) {
                    log.error(
                        "连接远程失败,conId:{}, clientPort:{}, fromHost:{},fromPost:{},remoteHost:{}, remotePort:{}, failMsg:{}",
                        myId,
                        connectBean.getClientPort(),
                        fromHost, fromPort,
                        connectBean.getHost(), connectBean.getPort(),
                        conRs.cause().getMessage());
                    //连接不成功
                    netSocket.close();
                    return;
                }
                log.info("建立新的转发连接,conId:{},clientPort:{}, fromHost:{},fromPost:{},remoteHost:{}, remotePort:{}",
                    myId,
                    connectBean.getClientPort(),
                    fromHost, fromPort,
                    connectBean.getHost(), connectBean.getPort());



                NetSocket mainTargetNetSocket = conRs.result();
                List<NetSocket> copyConnectedNetSockets = new ArrayList<>();
                if (connectBean.getCopyTo() != null && connectBean.getCopyTo().size() > 0) {
                    AtomicInteger connectingNum = new AtomicInteger(connectBean.getCopyTo().size());
                    // 连接每个拷贝目标地址后再开启pump
                    connectBean.getCopyTo().forEach(targetAddr -> connectToCopyAddr(myId,
                        netSocket, mainTargetNetSocket, copyConnectedNetSockets, connectingNum, targetAddr));
                } else {
                    // 不需要拷贝，直接开启pump
                    startPump(myId, netSocket, mainTargetNetSocket, copyConnectedNetSockets);
                }


            });
        }).listen(connectBean.getClientPort(), connectBean.getListenHost());
    }

    @Override
    public void stop() throws Exception {
        log.info("关闭监听, listenHost:{}, listenPort:{}, toHost:{}, toPort:{}",
            connectBean.getListenHost(), connectBean.getClientPort(),
            connectBean.getHost(), connectBean.getPort());
        netServer.close();
        super.stop();
    }
}
